import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

public class Solution560 {

    private int[] getPrefixArr(int[] nums){
        int[] prefixArr = new int[nums.length + 1];
        int sum = 0;
        for (int i = 0; i < nums.length; i++) {
            sum += nums[i];
            prefixArr[i + 1] = sum;
        }
        return prefixArr;
    }

    public int subarraySum(int[] nums, int k) {
        int count = 0;
        int[] prefixArr = getPrefixArr(nums);
        Map<Integer, Integer> prefixCounter = new HashMap<>();
        for (int prefix : prefixArr) {
            count += prefixCounter.getOrDefault(prefix - k, 0);
            prefixCounter.put(prefix, prefixCounter.getOrDefault(prefix, 0) + 1);
        }
        return count;
    }

    public static void main(String[] args) {
        Solution560 s = new Solution560();
        s.subarraySum(new int[]{1,1,1}, 2);
    }
}
